package fos

import scala.collection.immutable.{ Set, ListSet }

abstract class Type {
  override def toString() = this match {
    case TypeVar(a) => a
    case TypeFun(a, b) => "(" + a + " -> " + b + ")"
    case TypeNat => "Nat"
    case TypeBool => "Bool"
  }
}

case class TypeVar(name: String) extends Type

/** Type Schemes are not types. */
case class TypeScheme(args: List[TypeVar], tp: Type) {
  def instantiate: Type = {
    def subs(ty: Type, oldVar: TypeVar, newVar: TypeVar): Type = ty match {
      case x: TypeVar if (x == oldVar) => newVar
      case TypeFun(a, b) => TypeFun(subs(a, oldVar, newVar), subs(b, oldVar, newVar))
      case _ => ty
    }
    (tp /: args)(subs(_, _, Names.getFresh))
  }

  override def toString = args.mkString("[", ", ", "].") + tp
}

object Names {
  var num: Int = 0
  def getFresh(): TypeVar = {
    num = num + 1
    new TypeVar("$x" + num)
  }
}

case class TypeFun(a: Type, b: Type) extends Type

object Type
object TypeNat extends Type
object TypeBool extends Type
object TypeNone extends Type

abstract class Substitution extends (Type => Type) {

  def lookup(t: TypeVar): Option[Type];

  //  var indent = 0

  def apply(tp: Type): Type = {
    //    println("  " * indent + "in: " + tp + "   subst: " + this)
    //    indent = indent + 1
    val result = tp match {
      case x: TypeVar => lookup(x) match { case Some(t) => this(t); case None => tp }
      case TypeFun(t1, t2) => TypeFun(apply(t1), apply(t2))
      case _ => tp
    }
    //    indent = indent - 1
    //    println("  " * indent + "out: " + result + "   subst: " + this)
    result
  }

  def apply(p: (Type, Type)): (Type, Type) = p match {
    case Pair(t1, t2) => (this(t1), this(t2))
  }

  def apply(env: List[(String, TypeScheme)]): List[(String, TypeScheme)] =
    env map { (pair) => (pair._1, TypeScheme(pair._2.args, apply(pair._2.tp))) }

  def compose(other: Substitution): Substitution = other match {
    case EmptySubst => this
    case _ => new ComposedSubst(this, other)
  }

  def findConflict(tv: TypeVar, ty: Type): Option[Type];

  override def toString() = toList.toString();
  def toList: List[(TypeVar, Type)];
}

class SimpleSubst(v: TypeVar, t: Type) extends Substitution {

  override def lookup(ty: TypeVar): Option[Type] = if (v == ty) Some(t) else None

  def findConflict(tv: TypeVar, ty: Type): Option[Type] = {
    def hasFunConflict(f1: TypeFun, f2: TypeFun): Boolean = {
      val TypeFun(a, b) = f1
      val TypeFun(c, d) = f2
      (a, b, c, d) match {
        case (_: TypeVar, _: TypeVar, _, _) | (_: TypeVar, _, _, _: TypeVar) | (_, _: TypeVar, _: TypeVar, _) | (_, _, _: TypeVar, _: TypeVar) => false
        case (a: TypeFun, b: TypeFun, c: TypeFun, d: TypeFun) => hasFunConflict(a, c) || hasFunConflict(b, d)
        case (a: TypeFun, _, c: TypeFun, _) => hasFunConflict(a, c)
        case (_, b: TypeFun, _, d: TypeFun) => hasFunConflict(b, d)
        case (_: TypeVar, _, _, _) | (_, _, _: TypeVar, _) => b != d
        case (_, _: TypeVar, _, _) | (_, _, _, _: TypeVar) => a != c
        case (a, b, c, d) => !(a == b && c == d)
      }
    }

    (v, t, ty) match {
      case (`tv`, _, `t`) => None
      case (`tv`, t: TypeFun, ty: TypeFun) => if (hasFunConflict(t, ty)) Some(t) else None
      case (`tv`, _: TypeVar, _) | (`tv`, _, _: TypeVar) => None
      case (`tv`, t, _) => Some(t)
      case _ => None
    }
  }

  val toList: List[(TypeVar, Type)] = List((v, t))
}

class ComposedSubst(s1: Substitution, s2: Substitution) extends Substitution {
  def lookup(ty: TypeVar): Option[Type] = s2.lookup(ty) match {
    case None => s1.lookup(ty)
    case t => t
  }

  def findConflict(tv: TypeVar, ty: Type): Option[Type] = s1.findConflict(tv, ty) match {
    case None => s2.findConflict(tv, ty)
    case conflict => conflict
  }

  val toList: List[(TypeVar, Type)] = s2.toList ::: s1.toList
}

/** The empty substitution. */
object EmptySubst extends Substitution {
  override def lookup(t: TypeVar) = None
  override def compose(other: Substitution): Substitution = other
  override def toString: String = "No substitutions"
  val toList: List[(TypeVar, Type)] = Nil
  def findConflict(tv: TypeVar, ty: Type): Option[Type] = None
}
